#!/usr/bin/env python3
import os
import sys
import time
import threading
from collections import Counter
from typing import Optional, Dict, Any, Tuple, List

import numpy as np
import rospy
from std_msgs.msg import String


def _ensure_paths() -> None:
    """Ensure we can import detection pipeline and GStreamer wrapper from repo paths."""
    # scripts -> dog_motion -> src -> nav_ws -> repo_root
    repo_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
    yolo_detect_dir = os.path.join(repo_root, 'yolo-jetson', 'detect')
    yolo_rs2_dir = os.path.join(repo_root, 'yolo-jetson', 'rs2')
    if yolo_detect_dir not in sys.path:
        sys.path.append(yolo_detect_dir)
    if yolo_rs2_dir not in sys.path:
        sys.path.append(yolo_rs2_dir)


_ensure_paths()

from yolo_pipeline_optimized_cone_filtered import YOLOTrafficLightPipelineOptimized  # type: ignore
from GStreamerWrapper import GStreamerWrapper  # type: ignore


def _find_model_path() -> Optional[str]:
    """Find a default model path, preferring TensorRT .engine then .pt, under detect/weights."""
    # repo_root/yolo-jetson/detect/weights
    repo_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
    weights_dir = os.path.join(repo_root, 'yolo-jetson', 'detect', 'weights')
    candidates = [
        # Prefer engine
        os.path.join(weights_dir, 'all_yolo_all_epochs_0811', 'best.engine'),
        os.path.join(weights_dir, 'all_yolo_all_epochs_0808', 'best.engine'),
        os.path.join(weights_dir, 'all_yolo_all_epochs_0805', 'best.engine'),
        os.path.join(weights_dir, 'all_yolo_all_epochs_0804', 'best.engine'),
        # Fallback to pt
        os.path.join(weights_dir, 'all_yolo_all_epochs_0811', 'best.pt'),
        os.path.join(weights_dir, 'all_yolo_all_epochs_0808', 'best.pt'),
        os.path.join(weights_dir, 'all_yolo_all_epochs_0805', 'best.pt'),
        os.path.join(weights_dir, 'all_yolo_all_epochs_0804', 'best.pt'),
    ]
    for p in candidates:
        if os.path.exists(p):
            return p
    return None


class YoloGstConeFilteredNode:
    def __init__(self) -> None:
        # Params
        self.cache_size: int = int(rospy.get_param('~cache_size', 5))
        self.gst_width: int = int(rospy.get_param('~gst_width', 1280))
        self.gst_height: int = int(rospy.get_param('~gst_height', 720))
        self.device: Optional[str] = rospy.get_param('~device', None)
        self.use_fresh_pipeline_per_frame: bool = bool(rospy.get_param('~fresh_per_frame', True))

        # Model path (supports .engine or .pt)
        model_path_param = rospy.get_param('~model_path', '')
        self.model_path: Optional[str] = model_path_param if model_path_param else _find_model_path()
        if not self.model_path or not os.path.exists(self.model_path):
            rospy.logerr('Model weights not found. Please set ~model_path or place weights under yolo-jetson/detect/weights/.')
            raise RuntimeError('Model weights not found')

        # Preload once to validate model
        test_pipeline = YOLOTrafficLightPipelineOptimized(self.model_path, device=self.device, debug=False)
        del test_pipeline
        rospy.loginfo(f'Optimized model loaded: {self.model_path}')

        # ROS pub/sub
        self.result_pub = rospy.Publisher('/yolo_result', String, queue_size=10)
        self.cmd_sub = rospy.Subscriber('/yolo_detect', String, self._cmd_callback, queue_size=10)

        # State
        self.active_lock = threading.Lock()
        self.active: bool = False
        self.stop_event = threading.Event()
        self.worker_thread: Optional[threading.Thread] = None
        self.result_cache: List[Tuple[str, str, str]] = []

        # Optional persistent pipeline
        self.persistent_pipeline: Optional[YOLOTrafficLightPipelineOptimized] = None
        if not self.use_fresh_pipeline_per_frame:
            self.persistent_pipeline = YOLOTrafficLightPipelineOptimized(self.model_path, device=self.device, debug=False)

    def _cmd_callback(self, msg: String) -> None:
        data = (msg.data or '').strip().lower()
        if data == 'start':
            with self.active_lock:
                if self.active:
                    rospy.loginfo('Start requested but detection is already active. Ignoring.')
                    return
                self.active = True
                self.stop_event.clear()
                self.result_cache.clear()
                self.worker_thread = threading.Thread(target=self._run_detection_loop, daemon=True)
                self.worker_thread.start()
                rospy.loginfo('Detection started.')
        elif data == 'stop':
            with self.active_lock:
                if self.active:
                    rospy.loginfo('Stop requested. Stopping detection...')
                    self.stop_event.set()
                else:
                    rospy.loginfo('Stop requested but detection is not active.')

    def _process_single_frame(self, frame: np.ndarray, frame_idx: int) -> Dict[str, Any]:
        try:
            if self.use_fresh_pipeline_per_frame:
                pipeline = YOLOTrafficLightPipelineOptimized(self.model_path, device=self.device, debug=False)
            else:
                pipeline = self.persistent_pipeline

            results = pipeline.process_image(frame)

            if self.use_fresh_pipeline_per_frame:
                del pipeline

            return results
        except Exception as e:
            rospy.logwarn(f'Error processing frame {frame_idx}: {e}')
            return {
                'final_results': [],
                'cone_color': 'None',
                'number': 'None',
                'dashboard_state': 'None',
                'error': str(e),
                'processing_info': {},
            }

    def _majority_vote(self, items: List[str]) -> str:
        if not items:
            return 'None'
        c = Counter(items)
        return c.most_common(1)[0][0]

    def _publish_and_stop(self) -> None:
        cones = [c for c, _, _ in self.result_cache]
        numbers = [n for _, n, _ in self.result_cache]
        dashboards = [d for _, _, d in self.result_cache]

        cone_final = self._majority_vote(cones)
        number_final = self._majority_vote(numbers)
        dashboard_final = self._majority_vote(dashboards)

        result_str = f'{cone_final}-{number_final}-{dashboard_final}'
        rospy.loginfo(f'Publishing result: {result_str}')
        self.result_pub.publish(String(data=result_str))

        self.stop_event.set()
        self.result_cache.clear()

    def _run_detection_loop(self) -> None:
        gst = None
        try:
            rospy.loginfo('Initializing GStreamer source...')
            gst = GStreamerWrapper(width=self.gst_width, height=self.gst_height)
            frame_idx = 0

            while not rospy.is_shutdown() and not self.stop_event.is_set():
                frame = gst.GetFrame()
                if frame is None:
                    time.sleep(0.005)
                    continue

                results = self._process_single_frame(frame, frame_idx)

                cone_color = results.get('cone_color', 'None')
                number = results.get('number', 'None')
                dashboard_state = results.get('dashboard_state', 'None')

                rospy.loginfo(f'Frame {frame_idx}: cone={cone_color}, number={number}, dashboard={dashboard_state}')

                if cone_color != 'None' and number != 'None' and dashboard_state != 'None':
                    self.result_cache.append((str(cone_color), str(number), str(dashboard_state)))
                    rospy.loginfo(f'Cached valid result {len(self.result_cache)}/{self.cache_size}')
                    if len(self.result_cache) >= self.cache_size:
                        self._publish_and_stop()
                        break

                frame_idx += 1
        except Exception as e:
            rospy.logerr(f'Detection loop error: {e}')
        finally:
            try:
                if gst is not None:
                    gst.StopThread()
            except Exception:
                pass
            with self.active_lock:
                self.active = False
            rospy.loginfo('Detection stopped. Waiting for next start.')


def main() -> None:
    rospy.init_node('yolo_gst_cone_filtered_node', anonymous=False)
    node = YoloGstConeFilteredNode()
    rospy.loginfo('yolo_gst_cone_filtered_node ready. Send "start" on /yolo_detect to begin. Optional: send "stop" to cancel.')
    rospy.spin()


if __name__ == '__main__':
    main()


